#ifndef AVLTREE
#define AVLTREE

#include <stdlib.h>

//Code courtesy of Julienne Walker

#define MAX_LEVEL 32

//firs, enumeration for the balancing cases

typedef enum { leftDirMaxUnbalanced = -2, leftDirUnbalanced = -1, balanced = 0, rightDirUnbalanced = 1, rigthMaxUnbalanced = 2 } bFactor;
typedef enum { leftDir = 0, rightDir = 1 } direction;

//Node of the tree
template <class T>
struct Node {
   T data;						//stored data
   bFactor balance;				//balance value (-1, unbalanced on the leftDir, 0 balanced, 1 balanced on the rightDir)
   Node<T>* link[2];			//pointers to child (0 for leftDir child, 1 for rightDir child)
};

template <class T>
Node<T>* make_node( T val ) {
	Node<T>* res = new Node<T>;
	if ( res == NULL ) return res;
	res->data = val;
	res->balance = balanced;
	res->link[leftDir] = NULL;
	res->link[rightDir] = NULL;
	return res;
}

/**
 * AVL Tree template class
 * Implement an AVL tree with T data type.
 * It NEEDS that operators < and > are defined for the type T
 */
template <class T>
class AVLTree {

public:
	/**
	 * Class constructor: just put the root to null value.
	 */
	AVLTree( void ) { AVLroot = NULL; count = 0; }
	~AVLTree( void ) {}

	/**
	 * Check if the tree is empty
	 */
	bool empty( void ) {
		return (AVLroot == NULL);
	}

	/**
	 * Restore the tree to the original status
	 */
	void clear( void ) {
		//while( !empty() ) remove( minValue );
		AVLroot = NULL;
		count = 0;
	}

	/** 
	 * Insert element in the tree (Knuth's algorithm)
	 * Return true if the insertion goes well, false otherwise.
	 */
	bool insert ( T value )
	{
		/* Empty tree case */
		if ( AVLroot == NULL ) {
			AVLroot = make_node( value );
			minValue = value;
			maxValue = value;

			//allocation
			if ( AVLroot == NULL ) return false;
		}
		else {
			Node<T> head; /* False tree root */

			Node<T> *s, *t;     /* Place to rebalance and parent */
			Node<T> *p, *q;     /* Iterator and save pointer */
			direction dir;

			/* Set up false root to ease maintenance */
			t = &head;
			t->link[rightDir] = AVLroot;

			/* Search down the tree, saving rebalance points */
			for ( s = p = t->link[rightDir]; ; p = q ) {
				//no duplicates allowed
				if ( p->data == value ) return false;
				dir = static_cast<direction> (p->data < value);
				q = p->link[dir];

				if ( q == NULL ) break;

				if ( q->balance != 0 ) {
					t = p;
					s = q;
				}
			}
	 
			/* Insert the new node */
			//no duplicates allowed
			if ( p->data == value ) return false;
			p->link[dir] = q = make_node ( value );
			if ( q == NULL ) return false;
	 
			/* Update balance factors */
			for ( p = s; p != q; p = p->link[dir] ) {
				dir = static_cast<direction> (p->data < value);
				p->balance = static_cast<bFactor> ((static_cast<int> (p->balance) + ((dir == leftDir) ? -1 : +1)));
			}
	 
			q = s; /* Save rebalance point for parent fix */
	 
			/* Rebalance if necessary */
			if ( abs ( static_cast<int> (s->balance) ) > 1 ) {
				dir = static_cast<direction> (s->data < value);
				s = jsw_insert_balance ( s, dir );
			}
	 
			/* Fix parent */
			if ( q == head.link[rightDir] )
				AVLroot = s;
			else
				t->link[q == t->link[rightDir]] = s;
		}

		if ( value < minValue ) minValue = value;
		if ( value > maxValue ) maxValue = value;
	 
		count++;
		return true;
	}

	/**
	 * Find the specified value into the tree.
	 * Return a pointer to the element if it exists, null pointer otherwise
	 **/
	T* find( T value ) {
		//empty tree - return null
		if ( AVLroot == NULL ) 
			return NULL;

		Node<T> head; /* False tree root */

		Node<T> *s, *t;     /* Place to rebalance and parent */
		Node<T> *p, *q;     /* Iterator and save pointer */
		direction dir;

		/* Set up false root to ease maintenance */
		t = &head;
		t->link[rightDir] = AVLroot;

		/* Search down the tree, saving rebalance points */
		for ( s = p = t->link[rightDir]; ; p = q ) {
			//no duplicates allowed
			if ( p->data == value ) return  &(p->data);
			dir = static_cast<direction> (p->data < value);
			q = p->link[dir];

			if ( q == NULL )  return NULL;
		}

		return NULL;
	}

	/**
	 * Remove specified element from tree
	 * Return true if the elements gets removed, false if it's not present
	 */
	bool remove ( T value )
	{
		if ( AVLroot != NULL ) {

			Node<T> *it, *up[MAX_LEVEL];
			int upd[MAX_LEVEL], top = 0;
			int done = 0;

			it = AVLroot;

			for ( ; ; ) {
				/* Terminate if not found */
				if ( it == NULL ) 
					return false;
				else if ( it->data == value ) break;

				/* Push direction and node onto stack */
				upd[top] = it->data < value;
				up[top++] = it;

				it = it->link[upd[top - 1]];
			}

			/* Remove the node */
			if ( it->link[leftDir] == NULL || it->link[rightDir] == NULL ) {
				/* Which child is not null? */
				int dir = it->link[leftDir] == NULL;

				/* Fix parent */
				if ( top != 0 ) {
					up[top - 1]->link[upd[top - 1]] = it->link[dir];
					if ( it->data == minValue ) minValue = up[top-1]->data;
					if ( it->data == maxValue ) maxValue = up[top-1]->data;
				}
				else {
					AVLroot = it->link[dir];
					if ( AVLroot ) { minValue = AVLroot->data; maxValue = AVLroot->data; }
				}

				free ( it );
			}
			else {
				/* Find the inorder successor */
				Node<T> *heir = it->link[1];

				/* Save the path */
				upd[top] = 1;
				up[top++] = it;

				while ( heir->link[0] != NULL ) {
					upd[top] = 0;
					up[top++] = heir;
					heir = heir->link[0];
				}

				if ( it->data == minValue )
					minValue = heir->data;

				if ( it->data == maxValue )
					maxValue = heir->data;

				/* Swap data */
				it->data = heir->data;
				/* Unlink successor and fix parent */
				up[top - 1]->link[up[top - 1] == it] = heir->link[1];

				free ( heir );
			}

			/* Walk back up the search path */
			while ( --top >= 0 && !done ) {
				/* Update balance factors */
				up[top]->balance = static_cast<bFactor> (static_cast<int>(up[top]->balance) + (upd[top] != 0 ? -1 : +1));

				/* Terminate or rebalance as necessary */
				if ( abs ( up[top]->balance ) == 1 ) break;
				else if ( abs ( up[top]->balance ) > 1 ) {
					up[top] = jsw_remove_balance ( up[top], static_cast<direction>(upd[top]), &done );

					/* Fix parent */
					if ( top != 0 )
						up[top - 1]->link[upd[top - 1]] = up[top];
					else
						AVLroot = up[0];
				}
			}
		}

		return true;
	}


	/**
	 *  Return the minimum value of the tree
	 */
	T getMinValue( void ) {
		return minValue;
	}

	/**
	 *  Return the maximum value of the tree
	 */
	T getMaxValue( void ) {
		return maxValue;
	}

	/**
	 * Use carefully
	 */
	Node<T>* getRoot() { return AVLroot; }

private:
	Node<T>* AVLroot;

	T minValue;
	T maxValue;

	int count;
	/**
	 * Perform single rotation in a specified direction
	 */
	Node<T>* jsw_single ( Node<T>* root, direction dir )
	{
	   Node<T>* save = root->link[!dir];
	 
	   root->link[!dir] = save->link[dir];
	   save->link[dir] = root;
	 
	   return save;
	}

	/**
	 * Perform double rotation in a specified direction
	 */
	Node<T>* jsw_double ( Node<T>* root, direction dir )
	{
	  Node<T>* save = root->link[!dir]->link[dir];
	 
	  root->link[!dir]->link[dir] = save->link[!dir];
	  save->link[!dir] = root->link[!dir];
	  root->link[!dir] = save;

	  save = root->link[!dir];
	  root->link[!dir] = save->link[dir];
	  save->link[dir] = root;

	  return save;
	}

	/**
	 * Adjust tree balance
	 */
	void jsw_adjust_balance ( Node<T>* root, direction dir, bFactor bal )
	{
		Node<T> *n = root->link[dir];
		Node<T> *nn = n->link[!dir];

		if ( nn->balance == 0 )
			root->balance = n->balance = balanced;
		else if ( nn->balance == bal ) {
			root->balance = static_cast<bFactor> ( - static_cast<int>(bal) );
			n->balance = balanced;
		}
		else { /* nn->balance == -bal */
			root->balance = balanced;
			n->balance = bal;
		}

		nn->balance = balanced;
	}

	/**
	 * Insert element in the tree and balance
	 */
	Node<T> *jsw_insert_balance ( Node<T> *root, direction dir )
	{
		Node<T> *n = root->link[dir];
		bFactor bal = (dir == leftDir) ? leftDirUnbalanced : rightDirUnbalanced;

		if ( n->balance == bal ) {
			root->balance = n->balance = balanced;
			root = jsw_single ( root, static_cast<direction> (!(static_cast<bool>(dir))) );
		}
		else { /* n->balance == -bal */
			jsw_adjust_balance ( root, dir, bal );
			root = jsw_double ( root, static_cast<direction> (!(static_cast<bool>(dir))) );
		}

		return root;
	}

	
  /**
   * Remove element from the tree and balance
   */
	Node<T> *jsw_remove_balance ( Node<T> *root, direction dir, int *done )
	{
		Node<T> *n = root->link[!dir];
		bFactor bal = (dir == leftDir) ? leftDirUnbalanced : rightDirUnbalanced;

		if ( n->balance == static_cast<bFactor> ( - static_cast<int>(bal) ) ) {
			root->balance = n->balance = balanced;
			root = jsw_single ( root, dir );
		}
		else if ( n->balance == bal ) {
			jsw_adjust_balance ( root, static_cast<direction> (!(static_cast<bool>(dir))), static_cast<bFactor> ( - static_cast<int>(bal) ) );
			root = jsw_double ( root, dir );
		}
		else { /* n->balance == 0 */
			root->balance = static_cast<bFactor> ( - static_cast<int>(bal) );
			n->balance = bal;
			root = jsw_single ( root, dir );
			*done = 1;
		}

	return root;
	}
};

#endif //AVLTREE